import pandas as pd
from typing import Dict, Tuple, Optional
from datetime import datetime
import numpy as np
from loguru import logger


class BasisUtils:
    """基差计算工具类"""
    
    @staticmethod
    def calculate_basis(spot_price: float, futures_price: float, days_to_expiry: int) -> Tuple[float, float]:
        """
        计算基差点数和年化基差
        
        Args:
            spot_price: 现货价格
            futures_price: 期货价格  
            days_to_expiry: 距离到期天数
            
        Returns:
            Tuple[float, float]: (基差点数, 年化基差%)
        """
        try:
            # 基差点数 = 现货价格 - 期货价格
            basis_points = spot_price - futures_price
            
            # 年化基差 = (现货价格 - 期货价格) / 现货价格 * (365 / 剩余天数) * 100%
            if days_to_expiry <= 0 or spot_price <= 0:
                annualized_basis = 0.0
            else:
                annualized_basis = (spot_price - futures_price) / spot_price * (365 / days_to_expiry) * 100
                
            return basis_points, annualized_basis
            
        except Exception as e:
            logger.error(f"Error calculating basis: {e}")
            return 0.0, 0.0
    
    @staticmethod
    def add_basis_to_futures_data(spot_data: pd.DataFrame, futures_data: pd.DataFrame) -> pd.DataFrame:
        """
        将基差信息添加到期货数据中
        
        Args:
            spot_data: 现货数据DataFrame，需要包含 代码, 名称, 价格 列
            futures_data: 期货数据DataFrame，需要包含 合约, 最新价格, 距离到期日 列
            
        Returns:
            pd.DataFrame: 包含基差信息的期货数据DataFrame
        """
        try:
            # 创建现货代码到期货合约的映射关系
            index_mapping = {
                "000016": "IH",  # 上证50 -> IC
                "000300": "IF",  # 沪深300 -> IF  
                "000905": "IC",  # 中证500 -> IH
                "000852": "IM"   # 中证1000 -> IM
            }
            
            # 创建期货合约前缀到现货价格的映射
            futures_to_spot = {}
            for _, spot_row in spot_data.iterrows():
                spot_code = spot_row['代码']
                spot_price = float(spot_row['价格'])
                futures_prefix = index_mapping.get(spot_code)
                if futures_prefix:
                    futures_to_spot[futures_prefix] = spot_price
            
            # 复制期货数据并添加基差列
            enhanced_futures = futures_data.copy()
            enhanced_futures['基差点数'] = ""
            enhanced_futures['年化基差'] = ""
            
            # 为每个期货合约计算基差
            for idx, row in enhanced_futures.iterrows():
                contract_symbol = row['合约']
                futures_price = float(row['最新价格'])
                days_to_expiry = int(row['距离到期日'])
                
                # 找到对应的现货价格
                spot_price = None
                for prefix, price in futures_to_spot.items():
                    if contract_symbol.startswith(prefix):
                        spot_price = price
                        break
                
                if spot_price is not None:
                    # 计算基差
                    basis_points, annualized_basis = BasisUtils.calculate_basis(
                        spot_price, futures_price, days_to_expiry
                    )
                    enhanced_futures.at[idx, '基差点数'] = f"{basis_points:.2f}"
                    enhanced_futures.at[idx, '年化基差'] = f"{annualized_basis:.2f}%"
                else:
                    enhanced_futures.at[idx, '基差点数'] = "N/A"
                    enhanced_futures.at[idx, '年化基差'] = "N/A"
            
            return enhanced_futures
            
        except Exception as e:
            logger.error(f"Error adding basis to futures data: {e}")
            return futures_data 